{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to Create Custom Components\n",
    "\n",
    "This guide shows you how to create custom neurons, synapses, and other components in BrainPy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quick Start\n",
    "\n",
    "**Custom neuron template:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import brainpy\n",
    "import brainstate\n",
    "import brainunit as u\n",
    "import jax.numpy as jnp\n",
    "\n",
    "class CustomNeuron(brainpy.state.Neuron):\n",
    "    def __init__(self, size, **kwargs):\n",
    "        super().__init__(size, **kwargs)\n",
    "\n",
    "        # Parameters\n",
    "        self.tau = 10.0 * u.ms\n",
    "        self.V_th = -50.0 * u.mV\n",
    "\n",
    "        # States\n",
    "        self.V = brainstate.ShortTermState(jnp.zeros(size))\n",
    "        self.spike = brainstate.ShortTermState(jnp.zeros(size))\n",
    "\n",
    "    def reset_state(self, batch_size=None):\n",
    "        shape = self.size if batch_size is None else (batch_size, self.size)\n",
    "        self.V.value = jnp.zeros(shape)\n",
    "        self.spike.value = jnp.zeros(shape)\n",
    "\n",
    "    def update(self, x):\n",
    "        dt = brainstate.environ.get_dt()\n",
    "\n",
    "        # Dynamics\n",
    "        dV = -self.V.value / self.tau.to_decimal(u.ms) + x.to_decimal(u.nA)\n",
    "        self.V.value += dV * dt.to_decimal(u.ms)\n",
    "\n",
    "        # Spike generation\n",
    "        self.spike.value = (self.V.value >= self.V_th.to_decimal(u.mV)).astype(float)\n",
    "\n",
    "        # Reset\n",
    "        self.V.value = jnp.where(\n",
    "            self.spike.value > 0,\n",
    "            0.0,  # Reset voltage\n",
    "            self.V.value\n",
    "        )\n",
    "\n",
    "        return self.V.value\n",
    "\n",
    "    def get_spike(self):\n",
    "        return self.spike.value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom Neurons\n",
    "\n",
    "### Example 1: Adaptive LIF\n",
    "\n",
    "**LIF with spike-frequency adaptation:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AdaptiveLIF(brainpy.state.Neuron):\n",
    "    \"\"\"LIF neuron with adaptation current.\"\"\"\n",
    "\n",
    "    def __init__(self, size, tau=10*u.ms, tau_w=100*u.ms,\n",
    "                 V_th=-50*u.mV, V_reset=-65*u.mV, a=0.1*u.nA,\n",
    "                 b=0.5*u.nA, **kwargs):\n",
    "        super().__init__(size, **kwargs)\n",
    "\n",
    "        self.tau = tau\n",
    "        self.tau_w = tau_w\n",
    "        self.V_th = V_th\n",
    "        self.V_reset = V_reset\n",
    "        self.a = a  # Adaptation coupling\n",
    "        self.b = b  # Spike-triggered adaptation\n",
    "\n",
    "        # States\n",
    "        self.V = brainstate.ShortTermState(jnp.ones(size) * V_reset.to_decimal(u.mV))\n",
    "        self.w = brainstate.ShortTermState(jnp.zeros(size))  # Adaptation current\n",
    "        self.spike = brainstate.ShortTermState(jnp.zeros(size))\n",
    "\n",
    "    def reset_state(self, batch_size=None):\n",
    "        shape = self.size if batch_size is None else (batch_size, self.size)\n",
    "        self.V.value = jnp.ones(shape) * self.V_reset.to_decimal(u.mV)\n",
    "        self.w.value = jnp.zeros(shape)\n",
    "        self.spike.value = jnp.zeros(shape)\n",
    "\n",
    "    def update(self, I_ext):\n",
    "        dt = brainstate.environ.get_dt()\n",
    "\n",
    "        # Membrane potential dynamics\n",
    "        dV = (-self.V.value + self.V_reset.to_decimal(u.mV) + I_ext.to_decimal(u.nA) - self.w.value) / self.tau.to_decimal(u.ms)\n",
    "        self.V.value += dV * dt.to_decimal(u.ms)\n",
    "\n",
    "        # Adaptation dynamics\n",
    "        dw = (self.a.to_decimal(u.nA) * (self.V.value - self.V_reset.to_decimal(u.mV)) - self.w.value) / self.tau_w.to_decimal(u.ms)\n",
    "        self.w.value += dw * dt.to_decimal(u.ms)\n",
    "\n",
    "        # Spike generation\n",
    "        self.spike.value = (self.V.value >= self.V_th.to_decimal(u.mV)).astype(float)\n",
    "\n",
    "        # Reset and adaptation jump\n",
    "        self.V.value = jnp.where(\n",
    "            self.spike.value > 0,\n",
    "            self.V_reset.to_decimal(u.mV),\n",
    "            self.V.value\n",
    "        )\n",
    "        self.w.value += self.spike.value * self.b.to_decimal(u.nA)\n",
    "\n",
    "        return self.V.value\n",
    "\n",
    "    def get_spike(self):\n",
    "        return self.spike.value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 2: Izhikevich Neuron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Izhikevich(brainpy.state.Neuron):\n",
    "    \"\"\"Izhikevich neuron model.\"\"\"\n",
    "\n",
    "    def __init__(self, size, a=0.02, b=0.2, c=-65*u.mV, d=8*u.mV, **kwargs):\n",
    "        super().__init__(size, **kwargs)\n",
    "\n",
    "        self.a = a\n",
    "        self.b = b\n",
    "        self.c = c\n",
    "        self.d = d\n",
    "\n",
    "        # States\n",
    "        self.V = brainstate.ShortTermState(jnp.ones(size) * c.to_decimal(u.mV))\n",
    "        self.u = brainstate.ShortTermState(jnp.zeros(size))\n",
    "        self.spike = brainstate.ShortTermState(jnp.zeros(size))\n",
    "\n",
    "    def reset_state(self, batch_size=None):\n",
    "        shape = self.size if batch_size is None else (batch_size, self.size)\n",
    "        self.V.value = jnp.ones(shape) * self.c.to_decimal(u.mV)\n",
    "        self.u.value = jnp.zeros(shape)\n",
    "        self.spike.value = jnp.zeros(shape)\n",
    "\n",
    "    def update(self, I):\n",
    "        dt = brainstate.environ.get_dt()\n",
    "\n",
    "        # Izhikevich dynamics\n",
    "        dV = (0.04 * self.V.value**2 + 5 * self.V.value + 140 - self.u.value + I.to_decimal(u.nA))\n",
    "        du = self.a * (self.b * self.V.value - self.u.value)\n",
    "\n",
    "        self.V.value += dV * dt.to_decimal(u.ms)\n",
    "        self.u.value += du * dt.to_decimal(u.ms)\n",
    "\n",
    "        # Spike and reset\n",
    "        self.spike.value = (self.V.value >= 30).astype(float)\n",
    "        self.V.value = jnp.where(self.spike.value > 0, self.c.to_decimal(u.mV), self.V.value)\n",
    "        self.u.value = jnp.where(self.spike.value > 0, self.u.value + self.d.to_decimal(u.mV), self.u.value)\n",
    "\n",
    "        return self.V.value\n",
    "\n",
    "    def get_spike(self):\n",
    "        return self.spike.value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom Synapses\n",
    "\n",
    "### Example: Biexponential Synapse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BiexponentialSynapse(brainpy.state.Synapse):\n",
    "    \"\"\"Synapse with separate rise and decay.\"\"\"\n",
    "\n",
    "    def __init__(self, size, tau_rise=1*u.ms, tau_decay=5*u.ms, **kwargs):\n",
    "        super().__init__(size, **kwargs)\n",
    "\n",
    "        self.tau_rise = tau_rise\n",
    "        self.tau_decay = tau_decay\n",
    "\n",
    "        # States\n",
    "        self.h = brainstate.ShortTermState(jnp.zeros(size))  # Rising phase\n",
    "        self.g = brainstate.ShortTermState(jnp.zeros(size))  # Decaying phase\n",
    "\n",
    "    def reset_state(self, batch_size=None):\n",
    "        shape = self.size if batch_size is None else (batch_size, self.size)\n",
    "        self.h.value = jnp.zeros(shape)\n",
    "        self.g.value = jnp.zeros(shape)\n",
    "\n",
    "    def update(self, x):\n",
    "        dt = brainstate.environ.get_dt()\n",
    "\n",
    "        # Two-stage dynamics\n",
    "        dh = -self.h.value / self.tau_rise.to_decimal(u.ms) + x\n",
    "        dg = -self.g.value / self.tau_decay.to_decimal(u.ms) + self.h.value\n",
    "\n",
    "        self.h.value += dh * dt.to_decimal(u.ms)\n",
    "        self.g.value += dg * dt.to_decimal(u.ms)\n",
    "\n",
    "        return self.g.value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example: NMDA Synapse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NMDASynapse(brainpy.state.Synapse):\n",
    "    \"\"\"NMDA receptor with voltage dependence.\"\"\"\n",
    "\n",
    "    def __init__(self, size, tau=100*u.ms, a=0.5/u.mM, Mg=1.0*u.mM, **kwargs):\n",
    "        super().__init__(size, **kwargs)\n",
    "\n",
    "        self.tau = tau\n",
    "        self.a = a\n",
    "        self.Mg = Mg\n",
    "\n",
    "        self.g = brainstate.ShortTermState(jnp.zeros(size))\n",
    "\n",
    "    def reset_state(self, batch_size=None):\n",
    "        shape = self.size if batch_size is None else (batch_size, self.size)\n",
    "        self.g.value = jnp.zeros(shape)\n",
    "\n",
    "    def update(self, x, V_post=None):\n",
    "        \"\"\"Update with optional postsynaptic voltage.\"\"\"\n",
    "        dt = brainstate.environ.get_dt()\n",
    "\n",
    "        # Conductance dynamics\n",
    "        dg = -self.g.value / self.tau.to_decimal(u.ms) + x\n",
    "        self.g.value += dg * dt.to_decimal(u.ms)\n",
    "\n",
    "        # Voltage-dependent magnesium block\n",
    "        if V_post is not None:\n",
    "            mg_block = 1 / (1 + self.Mg.to_decimal(u.mM) * self.a.to_decimal(1/u.mM) * jnp.exp(-0.062 * V_post.to_decimal(u.mV)))\n",
    "            return self.g.value * mg_block\n",
    "        else:\n",
    "            return self.g.value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom Learning Rules\n",
    "\n",
    "### Example: Simplified STDP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleSTDP(brainstate.nn.Module):\n",
    "    \"\"\"Simplified STDP learning rule.\"\"\"\n",
    "\n",
    "    def __init__(self, n_pre, n_post, A_plus=0.01, A_minus=0.01,\n",
    "                 tau_plus=20*u.ms, tau_minus=20*u.ms):\n",
    "        super().__init__()\n",
    "\n",
    "        self.A_plus = A_plus\n",
    "        self.A_minus = A_minus\n",
    "        self.tau_plus = tau_plus\n",
    "        self.tau_minus = tau_minus\n",
    "\n",
    "        # Learnable weights\n",
    "        self.W = brainstate.ParamState(jnp.ones((n_pre, n_post)) * 0.5)\n",
    "\n",
    "        # Eligibility traces\n",
    "        self.pre_trace = brainstate.ShortTermState(jnp.zeros(n_pre))\n",
    "        self.post_trace = brainstate.ShortTermState(jnp.zeros(n_post))\n",
    "\n",
    "    def reset_state(self, batch_size=None):\n",
    "        shape_pre = self.W.value.shape[0] if batch_size is None else (batch_size, self.W.value.shape[0])\n",
    "        shape_post = self.W.value.shape[1] if batch_size is None else (batch_size, self.W.value.shape[1])\n",
    "        self.pre_trace.value = jnp.zeros(shape_pre)\n",
    "        self.post_trace.value = jnp.zeros(shape_post)\n",
    "\n",
    "    def update(self, pre_spike, post_spike):\n",
    "        dt = brainstate.environ.get_dt()\n",
    "\n",
    "        # Update traces\n",
    "        self.pre_trace.value += -self.pre_trace.value / self.tau_plus.to_decimal(u.ms) * dt.to_decimal(u.ms) + pre_spike\n",
    "        self.post_trace.value += -self.post_trace.value / self.tau_minus.to_decimal(u.ms) * dt.to_decimal(u.ms) + post_spike\n",
    "\n",
    "        # Weight updates\n",
    "        # LTP: pre spike finds existing post trace\n",
    "        dw_ltp = self.A_plus * jnp.outer(pre_spike, self.post_trace.value)\n",
    "\n",
    "        # LTD: post spike finds existing pre trace\n",
    "        dw_ltd = -self.A_minus * jnp.outer(self.pre_trace.value, post_spike)\n",
    "\n",
    "        # Update weights\n",
    "        self.W.value = jnp.clip(self.W.value + dw_ltp + dw_ltd, 0, 1)\n",
    "\n",
    "        return jnp.dot(pre_spike, self.W.value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom Network Architectures\n",
    "\n",
    "### Example: Liquid State Machine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LiquidStateMachine(brainstate.nn.Module):\n",
    "    \"\"\"Reservoir computing with spiking neurons.\"\"\"\n",
    "\n",
    "    def __init__(self, n_input=100, n_reservoir=1000, n_output=10):\n",
    "        super().__init__()\n",
    "\n",
    "        # Input projection (trainable)\n",
    "        self.input_weights = brainstate.ParamState(\n",
    "            brainstate.random.randn(n_input, n_reservoir) * 0.1\n",
    "        )\n",
    "\n",
    "        # Reservoir (fixed random recurrent network)\n",
    "        self.reservoir = brainpy.state.LIF(n_reservoir, V_rest=-65*u.mV, V_th=-50*u.mV, tau=10*u.ms)\n",
    "\n",
    "        # Fixed random recurrent weights\n",
    "        w_reservoir = brainstate.random.randn(n_reservoir, n_reservoir) * 0.01\n",
    "        mask = (brainstate.random.rand(n_reservoir, n_reservoir) < 0.1).astype(float)\n",
    "        self.reservoir_weights = w_reservoir * mask  # Not a ParamState (fixed)\n",
    "\n",
    "        # Readout (trainable)\n",
    "        self.readout = brainpy.state.Readout(n_reservoir, n_output)\n",
    "\n",
    "    def update(self, x):\n",
    "        # Input to reservoir\n",
    "        reservoir_input = jnp.dot(x, self.input_weights.value) * u.nA\n",
    "\n",
    "        # Reservoir recurrence\n",
    "        spk = self.reservoir.get_spike()\n",
    "        recurrent_input = jnp.dot(spk, self.reservoir_weights) * u.nA\n",
    "\n",
    "        # Update reservoir\n",
    "        self.reservoir(reservoir_input + recurrent_input)\n",
    "\n",
    "        # Readout from reservoir state\n",
    "        output = self.readout(self.reservoir.get_spike())\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom Input Encoders\n",
    "\n",
    "### Example: Temporal Contrast Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TemporalContrastEncoder(brainstate.nn.Module):\n",
    "    \"\"\"Encode images as spike timing based on contrast.\"\"\"\n",
    "\n",
    "    def __init__(self, n_pixels, max_time=100, threshold=0.1):\n",
    "        super().__init__()\n",
    "        self.n_pixels = n_pixels\n",
    "        self.max_time = max_time\n",
    "        self.threshold = threshold\n",
    "\n",
    "    def encode(self, image):\n",
    "        \"\"\"Convert image to spike timing.\n",
    "\n",
    "        Args:\n",
    "            image: Array of pixel values [0, 1]\n",
    "\n",
    "        Returns:\n",
    "            spike_times: When each pixel spikes (or max_time if no spike)\n",
    "        \"\"\"\n",
    "        # Higher intensity → earlier spike\n",
    "        spike_times = jnp.where(\n",
    "            image > self.threshold,\n",
    "            self.max_time * (1 - image),  # Invert: bright pixels spike early\n",
    "            self.max_time  # Below threshold: no spike\n",
    "        )\n",
    "\n",
    "        return spike_times\n",
    "\n",
    "    def decode_to_spikes(self, spike_times, current_time):\n",
    "        \"\"\"Get spikes at current simulation time.\"\"\"\n",
    "        spikes = (spike_times == current_time).astype(float)\n",
    "        return spikes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Best Practices\n",
    "\n",
    "✅ **Inherit from base classes**\n",
    "   - `brainpy.state.Neuron` for neurons\n",
    "   - `brainpy.state.Synapse` for synapses\n",
    "   - `brainstate.nn.Module` for general components\n",
    "\n",
    "✅ **Use ShortTermState for dynamics**\n",
    "   - Reset each trial\n",
    "   - Temporary variables\n",
    "\n",
    "✅ **Use ParamState for learnable parameters**\n",
    "   - Trained by optimizers\n",
    "   - Saved in checkpoints\n",
    "\n",
    "✅ **Implement reset_state()**\n",
    "   - Handle batch_size parameter\n",
    "   - Initialize all ShortTermStates\n",
    "\n",
    "✅ **Use physical units**\n",
    "   - All parameters with `brainunit`\n",
    "   - Convert for computation with `.to_decimal()`\n",
    "\n",
    "✅ **Follow naming conventions**\n",
    "   - `V` for voltage\n",
    "   - `spike` for spike indicator\n",
    "   - `g` for conductance\n",
    "   - `w` for weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing Custom Components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_custom_neuron():\n",
    "    \"\"\"Test custom neuron implementation.\"\"\"\n",
    "\n",
    "    neuron = CustomNeuron(size=10)\n",
    "    brainstate.nn.init_all_states(neuron)\n",
    "\n",
    "    # Test 1: Initialization\n",
    "    assert neuron.V.value.shape == (10,)\n",
    "    assert jnp.all(neuron.V.value == 0)\n",
    "\n",
    "    # Test 2: Response to input\n",
    "    strong_input = jnp.ones(10) * 10.0 * u.nA\n",
    "    for _ in range(100):\n",
    "        neuron(strong_input)\n",
    "\n",
    "    spike_count = jnp.sum(neuron.spike.value)\n",
    "    assert spike_count > 0, \"Neuron should spike with strong input\"\n",
    "\n",
    "    # Test 3: Batch dimension\n",
    "    brainstate.nn.init_all_states(neuron, batch_size=5)\n",
    "    assert neuron.V.value.shape == (5, 10)\n",
    "\n",
    "    print(\"✅ Custom neuron tests passed\")\n",
    "\n",
    "test_custom_neuron()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Complete Example\n",
    "\n",
    "**Putting it all together:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Custom components\n",
    "class MyNeuron(brainpy.state.Neuron):\n",
    "    # ... (see examples above)\n",
    "    pass\n",
    "\n",
    "class MySynapse(brainpy.state.Synapse):\n",
    "    # ... (see examples above)\n",
    "    pass\n",
    "\n",
    "# Use in network\n",
    "class CustomNetwork(brainstate.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.pre = MyNeuron(size=100)\n",
    "        self.post = MyNeuron(size=50)\n",
    "\n",
    "        self.projection = brainpy.state.AlignPostProj(\n",
    "            comm=brainstate.nn.EventFixedProb(100, 50, prob=0.1, weight=0.5*u.mS),\n",
    "            syn=MySynapse.desc(50),  # Use custom synapse\n",
    "            out=brainpy.state.CUBA.desc(),\n",
    "            post=self.post\n",
    "        )\n",
    "\n",
    "    def update(self, inp):\n",
    "        spk_pre = self.pre.get_spike()\n",
    "        self.projection(spk_pre)\n",
    "        self.pre(inp)\n",
    "        self.post(0*u.nA)\n",
    "        return self.post.get_spike()\n",
    "\n",
    "# Use network\n",
    "net = CustomNetwork()\n",
    "brainstate.nn.init_all_states(net)\n",
    "\n",
    "for _ in range(100):\n",
    "    output = net(input_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "**Component creation checklist:**\n",
    "\n",
    "```python\n",
    "✅ Inherit from brainpy.state.Neuron, brainpy.state.Synapse, or brainstate.nn.Module\n",
    "✅ Define __init__ with parameters\n",
    "✅ Create states (ShortTermState or ParamState)\n",
    "✅ Implement reset_state(batch_size=None)\n",
    "✅ Implement update() method\n",
    "✅ Use physical units throughout\n",
    "✅ Test with different batch sizes\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## See Also\n",
    "\n",
    "- Core Concepts: State Management\n",
    "- Core Concepts: Neurons\n",
    "- Core Concepts: Synapses\n",
    "- Tutorials: Synaptic Plasticity"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Ecosystem-py",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
